import warnings

import matplotlib.text as mtext
import matplotlib.axes as maxes
import matplotlib.transforms as mtransforms
import matplotlib.artist as artist
import matplotlib.axis as maxis
import matplotlib.ticker as mticker

import matplotlib
from matplotlib  import rcParams
#from matplotlib.collections import LineCollection
import matplotlib.lines as mlines
import matplotlib.artist as artist
import matplotlib.cbook as cbook

import pywcsgrid

import BaselineAlignedXticklabelsAxes

#from coords import FK5
#from coords.coord_system import coord_system_guess
try:
    from coords.coord_system import coord_system_guess, coord_sys
except ImportError:
    from pywcsgrid.wcs_helper import coord_system_guess, coord_sys


class NullDict(object):
    def __getitem__(self, item):
        return []

class DummyXAxis(maxis.XAxis):
    def __init__(self, axis):
        maxis.XAxis.__init__(self, axis.axes, axis.pickradius)
        self._scale = axis._scale #set_scale('linear')

    def set_ticks_onoff(self, onoff):

        ticks = list( self.get_major_ticks() )  # a copy
        ticks.extend( self.get_minor_ticks() )

        for t in ticks:
            t.tick1On = onoff
            t.tick2On = onoff


class DummyYAxis(maxis.YAxis):
    def __init__(self, axis):
        maxis.YAxis.__init__(self, axis.axes, axis.pickradius)
        self._scale = axis._scale #set_scale('linear')

    def set_ticks_onoff(self, onoff):

        ticks = list( self.get_major_ticks() )  # a copy
        ticks.extend( self.get_minor_ticks() )

        for t in ticks:
            t.tick1On = onoff
            t.tick2On = onoff



class TwinAxes(maxes.Axes):
    def __init__(self, ax):
        pass
        #maxes.Axes.__init__(self,
        #                    fig=ax.figure, rect=ax._position,
        #                    #sharex=ax, sharey=ax,
        #                    frameon=False)

        #self.xaxis = 
        #print 11
        #self.parent_axes = ax
        #self._position = ax._position
        
    def set_position(self, pos, which='both'):
        pass

    def apply_aspect(self, position=None):
        pass
    #self.set_position(self.parent_axes.get_position, "active")
        #pass


from mpl_transform import AuxTransformAxes, WcsSky2PixelTransform, \
                          WcsSky2SkyTransform
#AuxTransformAxes = pywcsgrid.mpl_transform.AuxTransformAxes
#    WcsSky2PixelTransform = pywcsgrid.mpl_transform.WcsSky2PixelTransform

    




class AxesCoordConvInterface(object):
    def __init__(self, *kl, **ka):

        #super(AxesCoordConvInterface, self).__init__()
        
        self._wcsgrid_wcsaxes = {}

        wcs = self.wcs
        self._wcsgrid_orig_coord_system = coord_system_guess(wcs.ctype[0],
                                                             wcs.ctype[1],
                                                             wcs.equinox)
        self._wcsgrid_target_coord_system = None


        ## add wcs axes
        ax = SubAxes(self,
                     WcsSky2PixelTransform(self.wcs))
        self._wcsgrid_wcsaxes[0] = ax


    def set_wcsgrid_coord_system(self, coord_system):
        self._wcsgrid_orig_coord_system = coord_system
        self.update_wcsgrid(self)


    def set_wcsgrid_target_coord_system(self, coord_system):
        if self._wcsgrid_orig_coord_system is None:
            raise ValueError("_pywcsgrid_orig_coord_system is not set")

        self._wcsgrid_target_coord_system = coord_system
        self.update_wcsgrid(self)
        
    #def draw(self, renderer):
    #   for k, ax in getattr(self, "_wcsgrid_wcsaxes", {}).items():
    #        ax.draw(renderer)
        #self._wcsgrid_wcsaxes[0].draw(renderer)

    def _get_artists(self):
        artists = []
        for k, ax in getattr(self, "_wcsgrid_wcsaxes", {}).items():
            artists.extend(ax._get_artists())
            #ax.draw(renderer)
        return artists
        #self._wcsgrid_wcsaxes[0].draw(renderer)


    def cla(self):

        # cla() is called before __init__ in maxes.Axes.__init__
        for k, ax in getattr(self, "_wcsgrid_wcsaxes", {}).items():
            ax.cla()


    def __getitem__(self, key):

        # check if key is a valid coord_sys instance
        if key != 0 and not isinstance(key, coord_sys):
            raise ValueError("argument key must be an instance of %s" % repr(coord_sys))

        if key not in self._wcsgrid_wcsaxes:
            if self._wcsgrid_orig_coord_system == key:
                self._wcsgrid_wcsaxes[key] = self._wcsgrid_wcsaxes[0]
            else:
                tr = WcsSky2SkyTransform(key,
                                         self._wcsgrid_orig_coord_system)
                #ax = AuxTransformAxes(self,
                #                      tr + self._wcsgrid_wcsaxes[0].transAux)
                ax = SubAxes(self,
                             tr + self._wcsgrid_wcsaxes[0].transAux)
                self._wcsgrid_wcsaxes[key] = ax

                
        return self._wcsgrid_wcsaxes[key]
            


class FixedFormatter(mticker.FixedFormatter):
    """Custom FixedFormatter. __call__ returns coordinate balue
    instead of '' when pos=None.  This gives some useful information
    for NavgationToolbar."""
    def __init__(self, seq):
        mticker.FixedFormatter.__init__(self, seq)
        
    def __call__(self, x, pos=None):
        'Return the format for tick val x at position pos'
        if pos is None:
            return "%g" % (x)
        elif pos>=len(self.seq): return ''
        else: return self.seq[pos]

    @classmethod
    def axis_set_ticklabels(cls, axis, ticklabels, *args, **kwargs):
        """
        Set the text values of the tick labels. Return a list of Text
        instances.  Use *kwarg* *minor=True* to select minor ticks.

        ACCEPTS: sequence of strings
        """

        self = axis
        
        # copied from axes.py
        #ticklabels = [str(l) for l in ticklabels]
        minor = kwargs.pop('minor', False)
        if minor:
            self.set_minor_formatter(cls(ticklabels))
            ticks = self.get_minor_ticks()
        else:
            self.set_major_formatter(cls(ticklabels) )
            ticks = self.get_major_ticks()

        self.set_major_formatter(cls(ticklabels) )

        ret = []
        for i, tick in enumerate(ticks):
            if i<len(ticklabels):
                tick.label1.set_text(ticklabels[i])
                ret.append(tick.label1)
            tick.label1.update(kwargs)
        return ret










class SubAxes(AuxTransformAxes):

    def __init__(self, parent_axes, aux_transform):
        self.__wcsgrid_parent_axes = parent_axes
        #self.__aux_grid_coords = []
        #self.__wcsgrid_aux_gridlines = []
        self.__wcsgrid_aux_gridlines = {}
        
        AuxTransformAxes.__init__(self, parent_axes, aux_transform)

    def set_wcsgrid_add_aux_grid_coord(self, lat, lon):
        #self.__aux_grid_coords.append((lat, lon))
        self.__wcsgrid_aux_gridlines[(lat, lon)] = None

        # TODO : we actually don't need a full update. Only gridlines
        # related this SubAxes are needed to.
        self.__wcsgrid_parent_axes.update_wcsgrid(self.__wcsgrid_parent_axes,
                                                  force_update=True)

    def get_wcsgrid_aux_gridlines(self):
        return self.__wcsgrid_aux_gridlines
    #return self.__aux_grid_coords, self.__wcsgrid_aux_gridlines

    def _get_artists(self):
        artists = []
        [artists.extend(l1) for l1 in self.__wcsgrid_aux_gridlines.itervalues()]
        artists.extend(self.lines)
        artists.extend(self.patches)
        artists.extend(self.texts)
        artists.extend(self.tables)
        artists.extend(self.artists)
        artists.extend(self.images) # images???
        # no legend. use parent's legend instead
        artists.extend(self.collections)
        
        return artists

    #def _set_wcsgrid_aux_gridlines(self, gridlines):
    #    self.__wcsgrid_aux_gridlines = gridlines

    #def get_wcsgrid_aux_gridlines(self):
    #    return self.__wcsgrid_aux_gridlines




import re
p = re.compile(r"\\u(\w*)\\d")
convert_table = {"h":r"\mathrm{h}",
                 "m":r"\mathrm{m}",
                 "s":r"\mathrm{s}",
                 "o":r"{\circ}",
                 "'":r"{\prime}",
                 '"':r"{\prime\prime}",
                 }
                     
                     
class Axes(BaselineAlignedXticklabelsAxes.AxesBase,
           maxes.Axes,
           AxesCoordConvInterface):
    
    def pgtext_to_latex(self, s):
        startindex=0
        r = []
        if not s: return ""
        for m in p.finditer(s):
            r.append(s[startindex:m.start()])
            c = m.groups()[0]
            r.append("^{" + convert_table.get(c, c) + "}\,")
            startindex=m.end()
        r.append(s[startindex:])
        r = "".join(r)
        r = r.replace("'", r"^{\prime}").replace('"',r"^{\prime\prime}")

        return "$"+r+"$"

    def _get_twin(self):
        ax2 = self.figure.add_axes(self.get_position(),
                                   #sharex=ax, sharey=ax, 
                                   frameon=False)

        #ax2._position = self._position
        #ax2.set_aspect(self.get_aspect())
        ax2.set_aspect("auto")

        #ax2.xaxis.set_ticks_position("none")
        #ax2.yaxis.set_ticks_position("none")
        ax2.xaxis.tick_top()
        ax2.yaxis.tick_right()
        #ax2.xaxis.set_label_position("top")
        #ax2.yaxis.set_label_position("right")
        return ax2
        
    def __init__(self, *kl, **ka):
        self.wcs = ka.pop("wcs")
        self.wcsgrid = None

        ka.setdefault("aspect", 1.)

        #self._gridOn = False
        self.wcsgrid_tickOn = True
        #self.wcsgrid_gridAxis2On = False

        self.wcsgrid_ticklines = []
        self.wcsgrid_gridlines = []
        self._wcsgrid_params = {}


        self.x1, self.x2, self.y1, self.y2 = None, None, None, None
        
        # we manually call __ini__ of each parent class because calling sequance is important

        BaselineAlignedXticklabelsAxes.AxesBase.__init__(self, *kl, **ka)
        maxes.Axes.__init__(self, *kl, **ka)
        AxesCoordConvInterface.__init__(self, *kl, **ka)
        #super(Axes, self).__init__(*kl, **ka)
        #super(Axes, self).__init__(*kl, **ka)
        
        #self._twin = TwinAxes(self) #self._get_twin()
        #self._twin.plot([1,2,3])

        
        #self.xaxis.set_ticks_position("none")
        #self.yaxis.set_ticks_position("none")
        self.xaxis.set_ticks_position("none")
        self.yaxis.set_ticks_position("none")

        self.xaxis_top = DummyXAxis(self.xaxis) #self._get_twin()
        self.yaxis_right = DummyYAxis(self.yaxis) #self._get_twin()
        #self.xaxis_top.set_ticks_position("none")
        #self.yaxis_right.set_ticks_position("none")

        self.xaxis_top.set_ticks_position("top")
        self.yaxis_right.set_ticks_position("right")
        self.xaxis_top.set_ticks_onoff(False)
        self.yaxis_right.set_ticks_onoff(False)
        
        self.set_wcsgrid_tick_size(None)
        
        self.set_wcsgrid_label_axis(bottom=1, left=2, top=None, right=None)

        self.set_wcsgrid_params(coord_format="GE",
                                label_density=(6, 6),
                                grid1=[], grid2=[])


        # wcsgrid is reevaluated whenever xlim and y lim changes. It is rather 
        # expensive calculation but couldn't fint any simple solution other than this.
        self.callbacks.connect("xlim_changed", self.update_wcsgrid)
        self.callbacks.connect("ylim_changed", self.update_wcsgrid)
        

    def set_wcsgrid_params(self, **ka):
        """
        coord_format="GE",
        label_density=(6, 6),
        grid1=[], grid2=[]):
        """
                          
        for k in ["coord_format", "label_density", "grid1", "grid2"]:
            if k in ka:
                self._wcsgrid_params[k] = ka.pop(k)
                
        if len(ka) > 0:
            raise ValueError("keyword name should be one of coord_format, label_densit, grid1, grid2")

        self.update_wcsgrid(self, True)
        
        
    @staticmethod
    def update_wcsgrid(self, force_update=False):
        # this needs to be declared as staticmethod to be connected to mpl events
        x1, x2 = self.get_xlim()
        y1, y2 = self.get_ylim()

        # 
        if force_update is False and (self.x1, self.x2, self.y1, self.y2) == (x1, x2, y1, y2):
            #update grids for SubAxes
            #self.update_wcsgrid_subaxes()
            return
        

        self.x1, self.x2, self.y1, self.y2 = x1, x2, y1, y2
        
        # update wcsgrid
        if self._wcsgrid_target_coord_system is not None:
            src_system = self._wcsgrid_orig_coord_system
            tgt_system = self._wcsgrid_target_coord_system
        else:
            src_system = None
            tgt_system = None

        self.wcsgrid = pywcsgrid.grid(self.wcs, x1, y1, x2, y2,
                                      src_system=src_system,
                                      dest_system=tgt_system, 
                                      **self._wcsgrid_params)


        # update gridlines
        self.wcsgrid_gridlines = self._update_gridlines(self.wcsgrid.grid_lines,
                                                        self.wcsgrid_gridlines)
        #### self.update_wcsgrid_gridlines()

        # update tics
        self.update_wcsgrid_ticks(self.wcsgrid_tick_size)

        # update tick labels
        self.update_wcsgrid_ticklabels()
        
        #update grids for SubAxes
        self.update_wcsgrid_subaxes()



    def update_wcsgrid_subaxes(self):

        # iterate over subaxes
        for coord_key, ax in self._wcsgrid_wcsaxes.items():

            # grid coords & gridlines 
            #coords, gridlines = ax.get_wcsgrid_aux_grid_coords_lines()

            # dictionary of {grid coords : gridlines}
            #coords, gridlines =
            coords_gridlines = ax.get_wcsgrid_aux_gridlines()
            
            if coord_key == 0:
                src_system=None
                dest_system=None
            else:
                src_system=self._wcsgrid_orig_coord_system
                dest_system=coord_key

            # update gridlines
            #for i in range(len(coords)):
            for k in coords_gridlines.iterkeys():
                lat1, lon1 = k
                if lat1 is None:
                    grid1, grid2 = None, [lon1]
                else:
                    grid1, grid2 = [lat1], None
                    
                _grid = pywcsgrid.grid(self.wcs,
                                       self.x1, self.y1, self.x2, self.y2,
                                       src_system=src_system,
                                       dest_system=dest_system,
                                       grid1=grid1, grid2=grid2)

                # _grid.grid_lines are lists of (x,y) tuples.
                
                # make mpl.Line2D instances from _grid.gridlines
                # copy line properties. Note both gridlines[i] and
                    
                _gridlines = self._update_gridlines(_grid.grid_lines,
                                                    coords_gridlines[k])
                #gridlines[i])

                # update gridlines. This updates subaxes.gridlines!
                coords_gridlines[k] = _gridlines
                #gridlines[i] = _gridlines


    

    def set_wcsgrid_tick_size(self, tick_size):
        self.wcsgrid_tick_size = tick_size
        if self.wcsgrid:
            self.update_wcsgrid_ticks(tick_size)
        
    def set_wcsgrid_ticks_toggle(self, b=None):
        if b is None:
            self.wcsgrid_tickOn = not self.wcsgrid_tickOn
        else:
            self.wcsgrid_tickOn = b

        self.update_wcsgrid_ticks(self.wcsgrid_tick_size)


        
    def update_wcsgrid_gridlines(self):
        self.wcsgrid_gridlines = self._update_gridlines(self.wcsgrid.grid_lines,
                                                        self.wcsgrid_gridlines)



    def _update_gridlines(self, gridlines, lines_for_copy_prop=None):
        l = []

        if lines_for_copy_prop:
            # remove lines first
            #[l1.remove() for l1 in self.wcsgrid_gridlines]
            line_for_copy_prop = lines_for_copy_prop[0]

            for x1, y1 in gridlines:
                l1 = mlines.Line2D(xdata=x1, ydata=y1)
                l1.update_from(line_for_copy_prop)
                l1.set_transform(self.transData) # may not need this
                self._set_artist_props(l1) # may not need this
                l1.set_clip_path(self.axesPatch)
                l.append(l1)
                #self.add_line(l1)

        else:
            for x1, y1 in gridlines:
                l1 = mlines.Line2D(xdata=x1, ydata=y1,
                                   color=rcParams['grid.color'],
                                   linestyle=rcParams['grid.linestyle'],
                                   linewidth=rcParams['grid.linewidth'],
                                   )
                l1.set_transform(self.transData)
                self._set_artist_props(l1)
                l1.set_clip_path(self.axesPatch)
                l.append(l1)
                    

        return cbook.silent_list('Line2D gridline', l)



    def update_wcsgrid_ticks(self, ts):

        if ts == None:
            ts = min(self.x2 - self.x1,
                     self.y2 - self.y1)/50.


        #[l.remove() for l in self.wcsgrid_ticklines]
            
        #from itertools import chain

        if self.wcsgrid_tickOn:
            lll = self.wcsgrid.get_ticklines(ts)
            self.wcsgrid_ticklines = self._get_tick_lines(lll)
        else:
            self.wcsgrid_ticklines = []
        #self.wcsgrid_ticklines = chain(self._get_xtick_lines(l[0]),
        #                               self._get_ytick_lines(l[1]))


    def _get_tick_lines(self, lll):

        l = []
        for x1, y1 in lll[0]:
            l1 = mlines.Line2D(xdata=x1, ydata=y1,
                               color=rcParams['xtick.color'],
                               linestyle="-",
                               linewidth=rcParams['grid.linewidth']*2.,
                               )
            l1.set_transform(self.transData)
            self._set_artist_props(l1)
            l1.set_clip_path(self.axesPatch)
            l.append(l1)
            

        for x1, y1 in lll[1]:
            l1 = mlines.Line2D(xdata=x1, ydata=y1,
                               color=rcParams['ytick.color'],
                               linestyle="-",
                               linewidth=rcParams['grid.linewidth']*2.,
                               )
            l1.set_transform(self.transData)
            self._set_artist_props(l1)
            l1.set_clip_path(self.axesPatch)
            l.append(l1)

        return cbook.silent_list('Line2D gridline', l)


                                                     




    def grid(self, b=None, **kwargs):
        if b is None:
            self._gridOn = not self._gridOn
        else:
            self._gridOn = b

        #self.update_wcsgrid_gridlines()

        if len(kwargs):
            for l1 in self.wcsgrid_gridlines:
                artist.setp(l1,**kwargs)






    def update_wcsgrid_ticklabels(self):
        get_axis = {1 : self.wcsgrid.axis1,
                    2 : self.wcsgrid.axis2,
                    None: NullDict()}

        # bottom
        axis = get_axis[self._wcsgrid_bottom_axis]
        self.set_xticks(axis["ticks_bottom"])
        #self.set_xticklabels(map(self.pgtext_to_latex,
        #                         axis["ticklabels_bottom"]))
        FixedFormatter.axis_set_ticklabels(self.xaxis,
                                           map(self.pgtext_to_latex,
                                               axis["ticklabels_bottom"]))

        axis = get_axis[self._wcsgrid_left_axis]
        self.set_yticks(axis["ticks_left"])
        #self.set_yticklabels(map(self.pgtext_to_latex,
        #                         axis["ticklabels_left"]))
        FixedFormatter.axis_set_ticklabels(self.yaxis,
                                           map(self.pgtext_to_latex,
                                               axis["ticklabels_left"]))

        axis = get_axis[self._wcsgrid_top_axis]
        self.xaxis_top.set_ticks(axis["ticks_top"])
        self.xaxis_top.set_ticklabels(map(self.pgtext_to_latex,
                                          axis["ticklabels_top"]))

        axis = get_axis[self._wcsgrid_right_axis]
        self.yaxis_right.set_ticks(axis["ticks_right"])
        self.yaxis_right.set_ticklabels(map(self.pgtext_to_latex,
                                            axis["ticklabels_right"]))


    
    def set_wcsgrid_label_axis(self, **ka):
        for k in ["bottom", "left", "top", "right"]:
            if k in ka:
                setattr(self, "_wcsgrid_%s_axis" % k, ka.pop(k))
        if len(ka) > 0:
            raise ValueError("keword name should be one of  bottom, left, top, right")
            
        if self.wcsgrid:
            self.update_wcsgrid_ticklabels()
            


    def get_images_and_artists(self):

        self.update_wcsgrid(self)
        
        images, artists = [], []

        images.extend(self.images)
        
        artists.extend(self.lines)
        artists.extend(self.patches)
        artists.extend(self.texts)
        artists.extend(self.tables)
        artists.extend(self.artists)
        #artists.extend(self.images) # images???
        # no legend. use parent's legend instead
        artists.extend(self.collections)


        artists.extend([self.xaxis, self.yaxis])
        artists.extend([self.xaxis_top, self.yaxis_right])

        artists.extend(self.wcsgrid_ticklines)

        if self._gridOn:
            artists.extend(self.wcsgrid_gridlines)


        artists.extend(AxesCoordConvInterface._get_artists(self))
                       #draw(self, renderer)

        
        return images, artists


    def draw(self, renderer):

        #self.update_wcsgrid(self)
        
        artists = []

        #self.xaxis_top.draw(renderer)
        #self.yaxis_right.draw(renderer)

        artists.extend([self.xaxis_top, self.yaxis_right])

        #for l1 in self.wcsgrid_ticklines:
        #    l1.draw(renderer)
        artists.extend(self.wcsgrid_ticklines)

        if self._gridOn:
            #for l1 in self.wcsgrid_gridlines:
            #    l1.draw(renderer)
            artists.extend(self.wcsgrid_gridlines)


        # subaxes grids
        #for coord_key, ax in self._wcsgrid_wcsaxes.items():
        #    coords, gridlines = ax.get_wcsgrid_aux_grid_coords_lines()
        #    [artists.extend(l1) for l1 in gridlines]


        artists.extend(AxesCoordConvInterface._get_artists(self))
                       #draw(self, renderer)


        # temporarily switch self.artists so that all artist are drawn
        # at one. This is preferred to preserve zorder of each artists
        orig_artists = self.artists
        artists.extend(orig_artists)
        self.artists = artists
        maxes.Axes.draw(self, renderer)
        self.artists = orig_artists






    def cla(self):
        self.wcsgrid_ticklines = []
        self.wcsgrid_gridlines = []
        self._wcsgrid_params = {}
        
        maxes.Axes.cla(self)
        AxesCoordConvInterface.cla(self)

    def imshow(self, *kl, **kw):

        origin = kw.get("origin", None)
        if origin is None: origin = rcParams['image.origin']
        
        if origin != "lower":
            warnings.warn("Warning : origin of the image is not 'lower' and this may cause incorrect drawing.")
        
        maxes.Axes.imshow(self, *kl, **kw)
        

        
Subplot = maxes.subplot_class_factory(Axes)
    
